package ginjwt

import (
    "crypto/rsa"
    "errors"
    "io/ioutil"
    "net/http"
    "strings"
    "time"

    "github.com/dgrijalva/jwt-go"
    "github.com/gin-gonic/gin"
)

type MapClaims map[string]interface{}

type GinJWTMiddleware struct {
    Realm string

    SigningAlgorithm string

    Key []byte

    Timeout time.Duration

    MaxRefresh time.Duration

    Authenticator func(c *gin.Context) (MapClaims, error)

    Authorizator func(claims MapClaims, c *gin.Context) bool

    PayloadFunc func(data interface{}) MapClaims

    Unauthorized func(*gin.Context, int, string)

    LoginResponse func(*gin.Context, int, string, time.Time)

    LogoutResponse func(*gin.Context, int)

    RefreshResponse func(*gin.Context, int, string, time.Time)

    IdentityHandler func(*gin.Context) interface{}

    IdentityKey string

    TokenLookup string

    TokenHeadName string

    TimeFunc func() time.Time

    HTTPStatusMessageFunc func(e error, c *gin.Context) string

    PrivKeyFile string

    PubKeyFile string

    privKey *rsa.PrivateKey

    pubKey *rsa.PublicKey

    SendCookie bool

    CookieMaxAge time.Duration

    SecureCookie bool

    CookieHTTPOnly bool

    CookieDomain string

    SendAuthorization bool

    DisabledAbort bool

    CookieName string

    CookieSameSite http.SameSite
}

var (
    ErrMissingSecretKey = errors.New("secret key is required")

    ErrForbidden = errors.New("you don't have permission to access this resource")

    ErrMissingAuthenticatorFunc = errors.New("ginJWTMiddleware.Authenticator func is undefined")

    ErrMissingLoginValues = errors.New("missing Username or Password")

    ErrFailedAuthentication = errors.New("incorrect Username or Password")

    ErrFailedTokenCreation = errors.New("failed to create JWT Token")

    ErrExpiredToken = errors.New("token is expired")

    ErrEmptyAuthHeader = errors.New("auth header is empty")

    ErrMissingExpField = errors.New("missing exp field")

    ErrWrongFormatOfExp = errors.New("exp must be float64 format")

    ErrInvalidAuthHeader = errors.New("auth header is invalid")

    ErrEmptyQueryToken = errors.New("query token is empty")

    ErrEmptyCookieToken = errors.New("cookie token is empty")

    ErrEmptyParamToken = errors.New("parameter token is empty")

    ErrInvalidSigningAlgorithm = errors.New("invalid signing algorithm")

    ErrNoPrivKeyFile = errors.New("private key file unreadable")

    ErrNoPubKeyFile = errors.New("public key file unreadable")

    ErrInvalidPrivKey = errors.New("private key invalid")

    ErrInvalidPubKey = errors.New("public key invalid")

    IdentityKey = "identity"
)

func New() *GinJWTMiddleware {
    x := &GinJWTMiddleware{}
    x.Realm = "WebServer"
    x.TokenLookup = "header:Authorization"
    x.SigningAlgorithm = "HS256"
    x.Timeout = time.Hour
    x.TimeFunc = time.Now
    x.TokenHeadName = "Bearer"
    x.Authorizator = func(data MapClaims, c *gin.Context) bool {
        return true
    }
    x.Unauthorized = func(c *gin.Context, code int, message string) {
        c.JSON(code, gin.H{
            "code":    code,
            "message": message,
        })
    }
    x.LoginResponse = func(c *gin.Context, code int, token string, expire time.Time) {
        c.JSON(http.StatusOK, gin.H{
            "code":   http.StatusOK,
            "token":  token,
            "expire": expire.Format(time.RFC3339),
        })
    }
    x.LogoutResponse = func(c *gin.Context, code int) {
        c.JSON(http.StatusOK, gin.H{
            "code": http.StatusOK,
        })
    }
    x.RefreshResponse = func(c *gin.Context, code int, token string, expire time.Time) {
        c.JSON(http.StatusOK, gin.H{
            "code":   http.StatusOK,
            "token":  token,
            "expire": expire.Format(time.RFC3339),
        })
    }
    x.IdentityKey = IdentityKey
    x.IdentityHandler = func(c *gin.Context) interface{} {
        claims := ExtractClaims(c)
        return claims[x.IdentityKey]
    }
    x.HTTPStatusMessageFunc = func(e error, c *gin.Context) string {
        return e.Error()
    }
    x.Realm = "gin jwt"
    x.CookieMaxAge = x.Timeout
    x.CookieName = "jwt"
    return x
}

func (s *GinJWTMiddleware) readKeys() error {
    err := s.privateKey()
    if err != nil {
        return err
    }
    err = s.publicKey()
    if err != nil {
        return err
    }
    return nil
}

func (s *GinJWTMiddleware) privateKey() error {
    keyData, err := ioutil.ReadFile(s.PrivKeyFile)
    if err != nil {
        return ErrNoPrivKeyFile
    }
    key, err := jwt.ParseRSAPrivateKeyFromPEM(keyData)
    if err != nil {
        return ErrInvalidPrivKey
    }
    s.privKey = key
    return nil
}

func (s *GinJWTMiddleware) publicKey() error {
    keyData, err := ioutil.ReadFile(s.PubKeyFile)
    if err != nil {
        return ErrNoPubKeyFile
    }
    key, err := jwt.ParseRSAPublicKeyFromPEM(keyData)
    if err != nil {
        return ErrInvalidPubKey
    }
    s.pubKey = key
    return nil
}

func (s *GinJWTMiddleware) usingPublicKeyAlgo() bool {
    switch s.SigningAlgorithm {
    case "RS256", "RS512", "RS384":
        return true
    }
    return false
}

func (s *GinJWTMiddleware) Setup() error {
    s.TokenHeadName = strings.TrimSpace(s.TokenHeadName)
    if len(s.TokenHeadName) == 0 {
        s.TokenHeadName = "Bearer"
    }

    if s.usingPublicKeyAlgo() {
        return s.readKeys()
    }

    if s.Key == nil {
        return ErrMissingSecretKey
    }
    return nil
}

func (s *GinJWTMiddleware) MiddlewareFunc() gin.HandlerFunc {
    return func(c *gin.Context) {
        s.middlewareImpl(c)
    }
}

func (s *GinJWTMiddleware) middlewareImpl(c *gin.Context) {
    claims, err := s.GetClaimsFromJWT(c)
    if err != nil {
        s.unauthorized(c, http.StatusUnauthorized, s.HTTPStatusMessageFunc(err, c))
        return
    }

    if claims["exp"] == nil {
        s.unauthorized(c, http.StatusBadRequest, s.HTTPStatusMessageFunc(ErrMissingExpField, c))
        return
    }

    if _, ok := claims["exp"].(float64); !ok {
        s.unauthorized(c, http.StatusBadRequest, s.HTTPStatusMessageFunc(ErrWrongFormatOfExp, c))
        return
    }

    if int64(claims["exp"].(float64)) < s.TimeFunc().Unix() {
        s.unauthorized(c, http.StatusUnauthorized, s.HTTPStatusMessageFunc(ErrExpiredToken, c))
        return
    }

    c.Set("JWT_PAYLOAD", claims)

    if !s.Authorizator(claims, c) {
        s.unauthorized(c, http.StatusForbidden, s.HTTPStatusMessageFunc(ErrForbidden, c))
        return
    }

    c.Next()
}

func (s *GinJWTMiddleware) GetClaimsFromJWT(c *gin.Context) (MapClaims, error) {
    token, err := s.ParseToken(c)

    if err != nil {
        return nil, err
    }

    if s.SendAuthorization {
        if v, ok := c.Get("JWT_TOKEN"); ok {
            c.Header("Authorization", s.TokenHeadName+" "+v.(string))
        }
    }

    claims := MapClaims{}
    for key, value := range token.Claims.(jwt.MapClaims) {
        claims[key] = value
    }

    return claims, nil
}

func (s *GinJWTMiddleware) LoginHandler(c *gin.Context) {
    if s.Authenticator == nil {
        s.unauthorized(c, http.StatusInternalServerError, s.HTTPStatusMessageFunc(ErrMissingAuthenticatorFunc, c))
        return
    }

    user_claims, err := s.Authenticator(c)
    if err != nil {
        s.unauthorized(c, http.StatusUnauthorized, s.HTTPStatusMessageFunc(err, c))
        return
    }

    token := jwt.New(jwt.GetSigningMethod(s.SigningAlgorithm))
    claims := token.Claims.(jwt.MapClaims)

    for key, value := range user_claims {
        claims[key] = value
    }

    expire := s.TimeFunc().Add(s.Timeout)
    claims["exp"] = expire.Unix()
    claims["orig_iat"] = s.TimeFunc().Unix()
    tokenString, err := s.signedString(token)

    if err != nil {
        s.unauthorized(c, http.StatusUnauthorized, s.HTTPStatusMessageFunc(ErrFailedTokenCreation, c))
        return
    }

    if s.SendCookie {
        expireCookie := s.TimeFunc().Add(s.CookieMaxAge)
        maxage := int(expireCookie.Unix() - s.TimeFunc().Unix())

        if s.CookieSameSite != 0 {
            c.SetSameSite(s.CookieSameSite)
        }

        c.SetCookie(
            s.CookieName,
            tokenString,
            maxage,
            "/",
            s.CookieDomain,
            s.SecureCookie,
            s.CookieHTTPOnly,
        )
    }

    s.LoginResponse(c, http.StatusOK, tokenString, expire)
}

func (s *GinJWTMiddleware) LogoutHandler(c *gin.Context) {

    if s.SendCookie {
        if s.CookieSameSite != 0 {
            c.SetSameSite(s.CookieSameSite)
        }

        c.SetCookie(
            s.CookieName,
            "",
            -1,
            "/",
            s.CookieDomain,
            s.SecureCookie,
            s.CookieHTTPOnly,
        )
    }

    s.LogoutResponse(c, http.StatusOK)
}

func (s *GinJWTMiddleware) signedString(token *jwt.Token) (string, error) {
    var tokenString string
    var err error
    if s.usingPublicKeyAlgo() {
        tokenString, err = token.SignedString(s.privKey)
    } else {
        tokenString, err = token.SignedString(s.Key)
    }
    return tokenString, err
}

func (s *GinJWTMiddleware) RefreshHandler(c *gin.Context) {
    tokenString, expire, err := s.RefreshToken(c)
    if err != nil {
        s.unauthorized(c, http.StatusUnauthorized, s.HTTPStatusMessageFunc(err, c))
        return
    }

    s.RefreshResponse(c, http.StatusOK, tokenString, expire)
}

func (s *GinJWTMiddleware) RefreshToken(c *gin.Context) (string, time.Time, error) {
    claims, err := s.CheckIfTokenExpire(c)
    if err != nil {
        return "", time.Now(), err
    }

    newToken := jwt.New(jwt.GetSigningMethod(s.SigningAlgorithm))
    newClaims := newToken.Claims.(jwt.MapClaims)

    for key := range claims {
        newClaims[key] = claims[key]
    }

    expire := s.TimeFunc().Add(s.Timeout)
    newClaims["exp"] = expire.Unix()
    newClaims["orig_iat"] = s.TimeFunc().Unix()
    tokenString, err := s.signedString(newToken)

    if err != nil {
        return "", time.Now(), err
    }

    if s.SendCookie {
        expireCookie := s.TimeFunc().Add(s.CookieMaxAge)
        maxage := int(expireCookie.Unix() - time.Now().Unix())

        if s.CookieSameSite != 0 {
            c.SetSameSite(s.CookieSameSite)
        }

        c.SetCookie(
            s.CookieName,
            tokenString,
            maxage,
            "/",
            s.CookieDomain,
            s.SecureCookie,
            s.CookieHTTPOnly,
        )
    }

    return tokenString, expire, nil
}

func (s *GinJWTMiddleware) CheckIfTokenExpire(c *gin.Context) (jwt.MapClaims, error) {
    token, err := s.ParseToken(c)
    if err != nil {

        validationErr, ok := err.(*jwt.ValidationError)
        if !ok || validationErr.Errors != jwt.ValidationErrorExpired {
            return nil, err
        }
    }

    claims := token.Claims.(jwt.MapClaims)

    origIat := int64(claims["orig_iat"].(float64))

    if origIat < s.TimeFunc().Add(-s.MaxRefresh).Unix() {
        return nil, ErrExpiredToken
    }

    return claims, nil
}

func (s *GinJWTMiddleware) TokenGenerator(data interface{}) (string, time.Time, error) {
    token := jwt.New(jwt.GetSigningMethod(s.SigningAlgorithm))
    claims := token.Claims.(jwt.MapClaims)

    if s.PayloadFunc != nil {
        for key, value := range s.PayloadFunc(data) {
            claims[key] = value
        }
    }

    expire := s.TimeFunc().UTC().Add(s.Timeout)
    claims["exp"] = expire.Unix()
    claims["orig_iat"] = s.TimeFunc().Unix()
    tokenString, err := s.signedString(token)
    if err != nil {
        return "", time.Time{}, err
    }

    return tokenString, expire, nil
}

func (s *GinJWTMiddleware) jwtFromHeader(c *gin.Context, key string) (string, error) {
    authHeader := c.Request.Header.Get(key)

    if authHeader == "" {
        return "", ErrEmptyAuthHeader
    }

    parts := strings.SplitN(authHeader, " ", 2)
    if !(len(parts) == 2 && parts[0] == s.TokenHeadName) {
        return "", ErrInvalidAuthHeader
    }

    return parts[1], nil
}

func (s *GinJWTMiddleware) jwtFromQuery(c *gin.Context, key string) (string, error) {
    token := c.Query(key)

    if token == "" {
        return "", ErrEmptyQueryToken
    }

    return token, nil
}

func (s *GinJWTMiddleware) jwtFromCookie(c *gin.Context, key string) (string, error) {
    cookie, _ := c.Cookie(key)

    if cookie == "" {
        return "", ErrEmptyCookieToken
    }

    return cookie, nil
}

func (s *GinJWTMiddleware) jwtFromParam(c *gin.Context, key string) (string, error) {
    token := c.Param(key)

    if token == "" {
        return "", ErrEmptyParamToken
    }

    return token, nil
}

func (s *GinJWTMiddleware) ParseToken(c *gin.Context) (*jwt.Token, error) {
    var token string
    var err error

    methods := strings.Split(s.TokenLookup, ",")
    for _, method := range methods {
        if len(token) > 0 {
            break
        }
        parts := strings.Split(strings.TrimSpace(method), ":")
        k := strings.TrimSpace(parts[0])
        v := strings.TrimSpace(parts[1])
        switch k {
        case "header":
            token, err = s.jwtFromHeader(c, v)
        case "query":
            token, err = s.jwtFromQuery(c, v)
        case "cookie":
            token, err = s.jwtFromCookie(c, v)
        case "param":
            token, err = s.jwtFromParam(c, v)
        }
    }

    if err != nil {
        return nil, err
    }

    return jwt.Parse(token, func(t *jwt.Token) (interface{}, error) {
        if jwt.GetSigningMethod(s.SigningAlgorithm) != t.Method {
            return nil, ErrInvalidSigningAlgorithm
        }
        if s.usingPublicKeyAlgo() {
            return s.pubKey, nil
        }

        c.Set("JWT_TOKEN", token)

        return s.Key, nil
    })
}

func (s *GinJWTMiddleware) ParseTokenString(token string) (*jwt.Token, error) {
    return jwt.Parse(token, func(t *jwt.Token) (interface{}, error) {
        if jwt.GetSigningMethod(s.SigningAlgorithm) != t.Method {
            return nil, ErrInvalidSigningAlgorithm
        }
        if s.usingPublicKeyAlgo() {
            return s.pubKey, nil
        }

        return s.Key, nil
    })
}

func (s *GinJWTMiddleware) unauthorized(c *gin.Context, code int, message string) {
    c.Header("WWW-Authenticate", "JWT realm="+s.Realm)
    if !s.DisabledAbort {
        c.Abort()
    }

    s.Unauthorized(c, code, message)
}

func ExtractClaims(c *gin.Context) MapClaims {
    claims, exists := c.Get("JWT_PAYLOAD")
    if !exists {
        return make(MapClaims)
    }

    return claims.(MapClaims)
}

func ExtractClaimsFromToken(token *jwt.Token) MapClaims {
    if token == nil {
        return make(MapClaims)
    }

    claims := MapClaims{}
    for key, value := range token.Claims.(jwt.MapClaims) {
        claims[key] = value
    }

    return claims
}

func GetToken(c *gin.Context) string {
    token, exists := c.Get("JWT_TOKEN")
    if !exists {
        return ""
    }

    return token.(string)
}
